# @Time : 2024/12/30 0:10
# @Author : ZHUYI
# @File : assessment
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, mean_squared_error
import matplotlib.pyplot as plt
import seaborn as sns

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 1. 读取数据
data = pd.read_csv('updated_pollution_dataset.csv')

# 2. 数据预处理
# 将目标变量 'Air Quality' 编码为数字
label_encoder = LabelEncoder()
data['Air Quality'] = label_encoder.fit_transform(data['Air Quality'])

# 特征变量和目标变量
X = data.drop(columns=['Air Quality'])
y = data['Air Quality']

# 3. 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 标准化特征
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 4. 线性回归模型
linear_model = LinearRegression()
linear_model.fit(X_train_scaled, y_train)
y_pred_linear = linear_model.predict(X_test_scaled)

# 输出线性回归模型的评估结果（均方误差）
mse = mean_squared_error(y_test, y_pred_linear)
print(f'线性回归模型的均方误差: {mse}')

# 5. 随机森林分类器
rf_model = RandomForestClassifier(random_state=42)
rf_model.fit(X_train, y_train)
y_pred_rf = rf_model.predict(X_test)

# 输出随机森林分类器的评估结果
print("随机森林分类器评估:")
print(confusion_matrix(y_test, y_pred_rf))
print(classification_report(y_test, y_pred_rf))

# 6. 可视化部分
# 6.1 随机森林混淆矩阵可视化
def plot_confusion_matrix(cm, model_name):
    plt.figure(figsize=(6, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False, xticklabels=label_encoder.classes_,
                yticklabels=label_encoder.classes_)
    plt.title(f'{model_name}的混淆矩阵')
    plt.xlabel('预测值')
    plt.ylabel('真实值')
    plt.show()

# 绘制随机森林混淆矩阵
cm_rf = confusion_matrix(y_test, y_pred_rf)
plot_confusion_matrix(cm_rf, '随机森林')

# 6.2 线性回归的预测值 vs 真实值
plt.figure(figsize=(8, 6))
plt.scatter(y_test, y_pred_linear, color='blue')
plt.plot([0, 3], [0, 3], color='red', linestyle='--')  # 理想情况下的预测线
plt.xlabel('真实值')
plt.ylabel('预测值')
plt.title('线性回归预测值与真实值对比')
plt.show()
